# -*- coding: utf-8 -*-
"""
Wargame AI的“上层AI” —— 用于比赛的第1阶段想定2010131194和2010141294
2020-06
"""
__author__ = 'Yu XH'

import copy
from myAI.hexmap.handle_map import *


class FactionAI:
    def __init__(self, myAgent):
        self.color = myAgent.color
        self.scenario = myAgent.scenario
        self.max_step = 1800
        self.task = {}  # 保存算子分配的任务信息{obj_id:['outputfire',param1,param2]}
        self.tactic = None  # 我方策略，'1-attack' | '2-defend'
        self.group_info = []  # 保存算子分组信息 [{group_id: 1, group_task: {ToOccupy:city}, operators: [opr1, opr2}]
        self.__left_side = self.color if self.scenario == 2010141294 else 1 - self.color  # 是否地图左边一方的判断，取值0|1
        self.__objs_mapping = self.__map_operators(myAgent)  # 根据算子的位置进行排序（从左上到右下排序），主要考虑到两个想定的特殊性，如此排序可以避免重复代码
        self.last_observation = myAgent.scenario_info[-1]  # 保存上一步状态，便于和当前状态比较

    def reset(self):
        self.color = None
        self.scenario = None
        self.max_step = 1800
        self.task = {}
        self.tactic = None
        self.group_info = None
        self.last_observation = None  # 保存上一步状态，便于和当前状态比较

    def update(self, myAgent):
        cur_step = myAgent.observation['time']['cur_step']
        our_oprs = myAgent.situations.our_oprs
        # 如果我方有算子死亡，则将其任务删除
        our_oprs_id = [opr['obj_id'] for opr in our_oprs]
        need_pop = [opr_id for opr_id in self.task.keys() if opr_id not in our_oprs_id]
        for pop_id in need_pop:
            self.task.pop(pop_id)

        if not our_oprs:
            self.task = {}
        elif self.ambush_1st_check(myAgent):  # 初始阶段，为所有算子下达指定地点的“伏击任务”
            for order, opr in enumerate(self.__objs_mapping):
                if opr['obj_id'] not in self.task.keys():
                    self.assign_ambush_task_1st(opr['obj_id'], order)
                elif self.task[opr['obj_id']][0] != 'ambush':
                    self.assign_ambush_task_1st(opr['obj_id'], order)
        elif self.control_condition_check(myAgent):  # 下达夺控任务
            # 1. 击毙敌人，直接夺控
            # 2. 我方得分劣势，在距离比赛结束20s内实施夺控
            self.assign_control_task(myAgent)
        elif self.ambush_2nd_check(myAgent):  # 第2次“伏击任务”
            # 1. 第1次伏击任务中，没有碰到敌人，在比赛500s后实施第2次伏击任务，将伏击点设置到夺控下附近
            self.assign_ambush_task_2nd(myAgent)
        elif self.anti_control_check(myAgent):  # 反夺控防守任务的判断
            # 1. 成功夺控后实施反夺控
            # 2. 我方得分无劣势，在距离比赛结束200s内实施反夺控
            self.assign_anti_control_task(myAgent)
        else:  # 以上不满足时实施追击任务
        # elif self.hunt_check(myAgent):  # 追击任务的判断
        #     # 1. 在夺控和反夺控任务中不实施追击任务
        #     # 2. 看见敌人时，实施追击任务
            self.assign_hunt_task(myAgent)

        # 更新上一次态势信息
        self.last_observation = myAgent.observation

    def assign_ambush_task_1st(self, obj_id, order):
        """分派第1次伏击任务"""
        if self.__left_side == 1:
            ambush_locs = [5457, 5557]
            self.task[obj_id] = ['ambush', ambush_locs[order]]
        else:
            ambush_locs = [5667, 5765]
            self.task[obj_id] = ['ambush', ambush_locs[order]]

    def assign_ambush_task_2nd(self, myAgent):
        our_oprs = myAgent.situations.our_oprs
        if len(our_oprs) == 0:
            self.task = {}
        if self.__left_side == 1:
            ambush_locs = [5559, 5460]
        else:
            ambush_locs = [5462, 5362]
        if len(our_oprs) == 1:
            self.task[our_oprs[0]['obj_id']] = ['ambush', ambush_locs[0]]
        else:
            self.task[our_oprs[0]['obj_id']] = ['ambush', ambush_locs[0]]
            self.task[our_oprs[1]['obj_id']] = ['ambush', ambush_locs[1]]

    def assign_control_task(self, myAgent):
        """分派夺控任务"""
        map_size = myAgent.get_map_size()
        city_loc = myAgent.observation['cities'][0]['coord']
        our_oprs = myAgent.situations.our_oprs
        if len(our_oprs) == 1:
            self.task[our_oprs[0]['obj_id']] = ['control', city_loc]
        else:
            control_task_flag = False
            for opr in our_oprs:
                if opr['obj_id'] in self.task.keys():
                    if self.task[opr['obj_id']][0] == 'control':
                        control_task_flag = True
            if self.no_empty_opr_check(our_oprs) and control_task_flag:
                pass
            else:
                city_loc_ind = mapid_2_matind(city_loc, map_size)
                cost_matrix = myAgent.my_map['cost_matrix'][0]
                dist_2_city = {}
                for opr in our_oprs:
                    opr_loc_ind = mapid_2_matind(opr['cur_hex'], map_size)
                    dist_2_city[opr['obj_id']] = cost_matrix[opr_loc_ind, city_loc_ind]
                if dist_2_city[our_oprs[0]['obj_id']] < dist_2_city[our_oprs[1]['obj_id']]:
                    control_obj, ambush_obj = our_oprs[0], our_oprs[1]
                elif dist_2_city[our_oprs[0]['obj_id']] > dist_2_city[our_oprs[1]['obj_id']]:
                    control_obj, ambush_obj = our_oprs[1], our_oprs[0]
                else:
                    if our_oprs[0]['blood'] <= our_oprs[1]['blood']:
                        control_obj, ambush_obj = our_oprs[0], our_oprs[1]
                    else:
                        control_obj, ambush_obj = our_oprs[1], our_oprs[0]
                self.task[control_obj['obj_id']] = ['control', city_loc]
                if self.__left_side == 1:
                    ambush_loc = 5460
                else:
                    ambush_loc = 5361
                self.task[ambush_obj['obj_id']] = ['ambush', ambush_loc]

    def assign_anti_control_task(self, myAgent):
        """夺控后，实施反夺控任务"""
        # 首先确定实施反夺控的算子
        city_loc = myAgent.observation['cities'][0]['coord']
        our_oprs = myAgent.situations.our_oprs
        anti_ctrl_opr_id = None
        for opr in our_oprs:
            if opr['obj_id'] in self.task.keys():
                if self.task[opr['obj_id']][0] in ['control', 'defend_firepower']:
                    anti_ctrl_opr_id = opr['obj_id']
                    break
        if not anti_ctrl_opr_id:
            if len(our_oprs) < 2:
                anti_ctrl_opr_id = our_oprs[0]['obj_id']
            else:
                map_size = myAgent.get_map_size()
                city_loc_ind = mapid_2_matind(city_loc, map_size)
                cost_matrix = myAgent.my_map['cost_matrix'][0]
                dist_2_city = {}
                for opr in our_oprs:
                    opr_loc_ind = mapid_2_matind(opr['cur_hex'], map_size)
                    dist_2_city[opr['obj_id']] = cost_matrix[opr_loc_ind, city_loc_ind]
                if dist_2_city[our_oprs[0]['obj_id']] <= dist_2_city[our_oprs[1]['obj_id']]:
                    anti_ctrl_opr_id = our_oprs[0]['obj_id']
                else:
                    anti_ctrl_opr_id = our_oprs[1]['obj_id']
        self.task[anti_ctrl_opr_id] = ['defend_firepower', city_loc, []]
        if self.__left_side == 1:
            ambush_loc = 5460
        else:
            ambush_loc = 5361
        for opr in our_oprs:
            if opr['obj_id'] != anti_ctrl_opr_id:
                if opr['obj_id'] not in self.task.keys():
                    self.task[opr['obj_id']] = ['ambush', ambush_loc]
                elif self.task[opr['obj_id']][0] != 'ambush':
                    self.task[opr['obj_id']] = ['ambush', ambush_loc]

    def assign_hunt_task(self, myAgent):
        """分派追击任务"""
        our_oprs = myAgent.situations.our_oprs
        enemies = myAgent.situations.enemies
        if len(enemies) == 1:
            tar_obj_id = enemies[0]['obj_id']
        else:
            n_our_oprs = len(our_oprs)
            our_locs = np.array([opr['cur_hex'] for opr in our_oprs])
            our_locs_ind = mapid_2_matind(our_locs, myAgent.get_map_size())
            dist_matrix = myAgent.my_map['dist_matrix'][our_locs_ind]
            avg_dist = []
            for enemy_info in enemies:
                avg_dist.append(np.sum(dist_matrix @ enemy_info["pos_probs"]) / n_our_oprs)
            min_item = np.argmin(np.array(avg_dist))
            tar_obj_id = enemies[min_item]['obj_id']
        for opr in our_oprs:
            self.task[opr['obj_id']] = ['hunt', tar_obj_id]

    def __map_operators(self, myAgent):
        objs_mapping = copy.deepcopy(myAgent.scenario_info[self.color]['operators'])
        objs_mapping.sort(key=lambda obj: obj['cur_hex'])
        return objs_mapping

    def ambush_1st_check(self, myAgent):
        """第1次伏击的判断"""
        cur_step = myAgent.observation['time']['cur_step']
        if cur_step >= 500:  # 第1次伏击的时间<500s
            return False
        our_oprs = myAgent.situations.our_oprs
        if len(our_oprs) < 2:
            return False
        for opr in our_oprs:  # 跟敌方没有接触
            for init_opr in myAgent.scenario_info[self.color]['operators']:
                if opr['obj_id'] == init_opr['obj_id']:
                    if opr['blood'] < init_opr['blood']:
                        return False
                    for key in opr['remain_bullet_nums'].keys():
                        if opr['remain_bullet_nums'][key] < init_opr['remain_bullet_nums'][key]:
                            return False
            if opr['see_enemy_bop_ids']:
                return False
        return True

    def control_condition_check(self, myAgent):
        """夺控条件的判断"""
        cur_step = myAgent.observation['time']['cur_step']
        our_oprs = myAgent.situations.our_oprs
        enemies = myAgent.situations.enemies
        if myAgent.observation['cities'][0]['flag'] == self.color:
            return False
        if not enemies:
            return True
        our_win_score = myAgent.observation['scores']['red_win'] if self.color == 0 else myAgent.observation['scores'][
            'blue_win']
        if our_win_score >= 0:
            return False
        cost_matrix = myAgent.my_map['cost_matrix'][0]
        map_size = myAgent.get_map_size()
        minimum_time_2_city = 10000
        city_loc = myAgent.observation['cities'][0]['coord']
        city_loc_ind = mapid_2_matind(city_loc, map_size)
        for opr in our_oprs:  # 更新到夺控点的最短距离
            opr_loc_ind = mapid_2_matind(opr['cur_hex'], map_size)
            dist_2_city = cost_matrix[opr_loc_ind, city_loc_ind]
            if dist_2_city < minimum_time_2_city:
                minimum_time_2_city = dist_2_city
        if our_win_score < 0 and (self.max_step - (cur_step + minimum_time_2_city) < 20):
            return True
        return False

    def ambush_2nd_check(self, myAgent):
        """第2次伏击的判断"""
        cur_step = myAgent.observation['time']['cur_step']
        if cur_step < 500:  # 第1次伏击的时间
            return False
        our_oprs = myAgent.situations.our_oprs
        if len(our_oprs) < 2:
            return False
        for opr in our_oprs:  # 跟敌方没有接触
            for init_opr in myAgent.scenario_info[self.color]['operators']:
                if opr['obj_id'] == init_opr['obj_id']:
                    if opr['blood'] < init_opr['blood']:
                        return False
                    for key in opr['remain_bullet_nums'].keys():
                        if opr['remain_bullet_nums'][key] < init_opr['remain_bullet_nums'][key]:
                            return False
            if opr['see_enemy_bop_ids']:
                return False
        return True

    def anti_control_check(self, myAgent):
        """反夺控任务的判断"""
        cur_city_flag = myAgent.observation['cities'][0]['flag']
        if cur_city_flag == self.color:
            return True
        our_win_score = myAgent.observation['scores']['red_win'] if self.color == 0 else myAgent.observation['scores'][
            'blue_win']
        cur_step = myAgent.observation['time']['cur_step']
        if our_win_score >= 0 and (self.max_step - cur_step < 200):
            return True
        return False

    def hunt_check(self, myAgent):
        """追击任务的判断"""
        if self.task.values():
            for task in self.task.values():
                if task[0] in ['control', 'defend_firepower']:
                    return False
        our_oprs = myAgent.situations.our_oprs
        enemies_in_sight = []
        for opr in our_oprs:
            enemies_in_sight += opr['see_enemy_bop_ids']
        if enemies_in_sight:  # 当前有敌人出现
            return True

    def no_empty_opr_check(self, our_oprs):
        """判断我方是否有空任务算子"""
        for opr in our_oprs:
            if opr['obj_id'] not in self.task.keys():
                return False
            if self.task[opr['obj_id']][0] in ['Failure', 'Success']:
                return False
        return True
